"""
词向量和数字向量的转换
"""
import pickle
from tqdm import tqdm

import config
from chatbot.word_sequence import WordSequence


def save_ws():
    # 创建WordSequence对象
    ws_q = WordSequence()
    ws_sim_q = WordSequence()

    # 读取文件内容
    with open(config.sort_q_data_path, mode="r", encoding="UTF-8") as f_q:
        q_lines = f_q.readlines()
    with open(config.sort_sim_q_data_path, mode="r", encoding="UTF-8") as f_sim_q:
        sim_q_lines = f_sim_q.readlines()

    # 对文件中的内容进行fit操作
    for q_line, sim_q_line in tqdm(zip(q_lines, sim_q_lines)):
        ws_q.fit(q_line.strip().split())
        ws_sim_q.fit(sim_q_line.strip().split())

    # 创建字典
    ws_q.build_vocab()
    ws_sim_q.build_vocab()

    # 保存模型
    pickle.dump(ws_q, open(config.sort_q_ws_model_path, mode="wb"))
    pickle.dump(ws_sim_q, open(config.sort_sim_q_ws_model_path, mode="wb"))



